#include "kernel_operator.h"
#include <type_traits>
#include "sinCompute.h"
#include "genericFresnelCos.h"
#include "genericFresnelAsymp.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 1;
class KernelFresnelCos {
public:
    __aicore__ inline KernelFresnelCos() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, 
                                int32_t loopCount, int32_t tileDataLenght, int32_t tailDataLenght, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->loopCount = loopCount;
        this->tileDataLenght = tileDataLenght;
        this->tailDataLenght = tailDataLenght;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x);
        yGm.SetGlobalBuffer((__gm__ DTYPE_X*)y);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataLenght * sizeof(DTYPE_X));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataLenght * sizeof(DTYPE_X));
        if constexpr (std::is_same_v<DTYPE_X, float>)
        {
          pipe.InitBuffer(QueueTmp, this->tileDataLenght * 7 * sizeof(float));  
          pipe.InitBuffer(QueueMask1, this->tileDataLenght / 8 * sizeof(uint8_t));
          pipe.InitBuffer(QueueMask2, this->tileDataLenght / 8 * sizeof(uint8_t));
        }
        else
        {
          pipe.InitBuffer(QueueTmp, this->tileDataLenght * 9 * sizeof(float));
          pipe.InitBuffer(QueueMask1, this->tileDataLenght / 8 * sizeof(uint8_t));
        }
    }
    __aicore__ inline void Process() {
      this->processDataLenght = this->tileDataLenght;
      for(int32_t i=0; i<this->loopCount; i++)
      {
        CopyIn(i);
        Compute(i);
        CopyOut(i);
      }
      this->processDataLenght = this->tailDataLenght;
      CopyIn(this->loopCount);
      Compute(this->loopCount);
      CopyOut(this->loopCount);
    }
 private:
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[progress * this->tileDataLenght], this->processDataLenght);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t progress) {
      LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
      LocalTensor<DTYPE_X> yLocal = outQueueY.AllocTensor<DTYPE_X>();
      auto tmp = QueueTmp.Get<float>();
      if constexpr (std::is_same_v<DTYPE_X, float>)
      {
        auto tmpOut1 = tmp;
        auto tmpOut2 = tmp[this->tileDataLenght];
        auto tmp1 = tmp[this->tileDataLenght*2];
        auto tmp2 = tmp[this->tileDataLenght*3];
        auto tmp3 = tmp[this->tileDataLenght*4];
        auto tmp4 = tmp[this->tileDataLenght*5];
        auto tmp5 = tmp[this->tileDataLenght*6];
        auto tmpMask1 = QueueMask1.Get<uint8_t>();
        auto tmpMask2 = QueueMask2.Get<uint8_t>();
        CompareScalar(tmpMask1, xLocal.ReinterpretCast<float>(), static_cast<float>(0), CMPMODE::LT, this->tileDataLenght);
        Abs(xLocal, xLocal, this->processDataLenght);
        CompareScalar(tmpMask2, xLocal.ReinterpretCast<float>(), static_cast<float>(1.6), CMPMODE::LT, this->tileDataLenght);
        GenericFresnelCos(tmpOut1, xLocal.ReinterpretCast<float>(), tmp5, yLocal.ReinterpretCast<float>(), this->processDataLenght);
        GenericFresnelAsymp(tmpOut2, xLocal.ReinterpretCast<float>(), tmp1, tmp2, tmp3, tmp4, tmp5, yLocal.ReinterpretCast<float>(), this->processDataLenght);
        Select(tmp1, tmpMask2, tmpOut1, tmpOut2, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataLenght);
        Muls(tmp2, tmp1, (float)(-1), this->processDataLenght);
        Select(yLocal.ReinterpretCast<float>(), tmpMask1, tmp2, tmp1, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataLenght);
      }
      else
      {
        auto tmpX = tmp;
        auto tmpOut1 = tmp[this->tileDataLenght];
        auto tmpOut2 = tmp[this->tileDataLenght*2];
        auto tmp1 = tmp[this->tileDataLenght*3];
        auto tmp2 = tmp[this->tileDataLenght*4];
        auto tmp3 = tmp[this->tileDataLenght*5];
        auto tmp4 = tmp[this->tileDataLenght*6];
        auto tmp5 = tmp[this->tileDataLenght*7];
        auto tmp6 = tmp[this->tileDataLenght*8];
        auto tmpMask1 = QueueMask1.Get<uint8_t>();
        ShiftRight(yLocal.ReinterpretCast<uint16_t>(), xLocal.ReinterpretCast<uint16_t>(), (uint16_t)15, this->processDataLenght);
        ShiftLeft(yLocal.ReinterpretCast<uint16_t>(), yLocal.ReinterpretCast<uint16_t>(), (uint16_t)15, this->processDataLenght);
        Abs(xLocal, xLocal, this->processDataLenght);
        Cast(tmpX, xLocal, RoundMode::CAST_NONE, this->processDataLenght);
        CompareScalar(tmpMask1, tmpX, static_cast<float>(1.6), CMPMODE::LT, this->tileDataLenght);
        GenericFresnelCos(tmpOut1, tmpX, tmp5, tmp6, this->processDataLenght);
        GenericFresnelAsymp(tmpOut2, tmpX, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, this->processDataLenght);
        Cast(tmp1.ReinterpretCast<half>(), tmpOut1, RoundMode::CAST_RINT, this->processDataLenght);
        Cast(tmp2.ReinterpretCast<half>(), tmpOut2, RoundMode::CAST_RINT, this->processDataLenght);
        Select(tmp3.ReinterpretCast<half>(), tmpMask1, tmp1.ReinterpretCast<half>(), tmp2.ReinterpretCast<half>(), SELMODE::VSEL_TENSOR_TENSOR_MODE, this->processDataLenght);
        Or(yLocal.ReinterpretCast<uint16_t>(), yLocal.ReinterpretCast<uint16_t>(), tmp3.ReinterpretCast<uint16_t>(), this->processDataLenght);
      }
      inQueueX.FreeTensor(xLocal);    
      outQueueY.EnQue<DTYPE_X>(yLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<DTYPE_X> yLocal = outQueueY.DeQue<DTYPE_X>();
        DataCopy(yGm[progress * this->tileDataLenght], yLocal, this->processDataLenght);
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp, QueueMask1, QueueMask2;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_X> yGm;
    int32_t loopCount;
    int32_t tileDataLenght;
    int32_t tailDataLenght;
    int32_t processDataLenght;
};
extern "C" __global__ __aicore__ void fresnel_cos(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    KernelFresnelCos op;
    op.Init(x, y, 
            tiling_data.loopCount, tiling_data.tileDataLenght, tiling_data.tailDataLenght, &pipe);
    op.Process();
}